# 截图宝 - 在线文档截图软件主窗体
# 作者： 魏新崇
# 邮箱： weixinchong [at] 163.com
# 项目主页：https://gitee.com/weixinchong/auto-igrab

import sys, pyautogui, os, time, numpy as np
from PyQt5.QtWidgets import QApplication, QWidget, QHBoxLayout, QGridLayout, QLabel, QPushButton, QSpinBox, QLineEdit, QDoubleSpinBox, QRadioButton, QProgressBar
from PyQt5.QtGui import QFont
from PIL import ImageGrab, Image
from win_base import Window_Base
from dispdialog import DisplacementDialog


class DocumentWindow(Window_Base):

    def __init__(self, parent = None):
        super(DocumentWindow, self).__init__('在线文档截图宝', (42, 96, 884, 975), parent)

        # 在基本空间的基础上添加控件
        # 1.翻页控制
        # 1.1 创建控件
        label_page_turn = QLabel('翻页控制')
        label_page_turn.setFont(self.font_title)
        label_method = QLabel('翻页策略：')
        self.method = 1
        self.rad_method1 = QRadioButton('点击下一页')
        self.rad_method2 = QRadioButton('更改页码')
        self.rad_method1.toggle()
        self.rad_method1.toggled.connect(self.update_method)
        self.rad_method2.toggled.connect(self.update_method)
        label_mouse_pos0 = QLabel('起始坐标：')
        self.spin_mouse_x0 = self.create_spin(450)
        self.spin_mouse_y0 = self.create_spin(848)
        btn_mouse_pos0_adjust = QPushButton('重定位')
        btn_mouse_pos0_adjust.clicked.connect(lambda: self.get_position(self.spin_mouse_x0, self.spin_mouse_y0))
        label_mouse_posp = QLabel('功能键坐标：')
        self.spin_mouse_xp = self.create_spin(575)
        self.spin_mouse_yp = self.create_spin(1007)
        btn_mouse_posp_adjust = QPushButton('重定位')
        btn_mouse_posp_adjust.clicked.connect(lambda: self.get_position(self.spin_mouse_xp, self.spin_mouse_yp))
        label_page_ydisp = QLabel('页垂直位移：')
        self.spin_page_ydisp = self.create_spin(0, vmin = -10000)
        btn_page_ydisp_adjust = QPushButton('确定位移')
        btn_page_ydisp_adjust.clicked.connect(self.get_ydisp)
        label_turn_interval = QLabel('翻页间隔：')
        self.spin_turn_interval = QDoubleSpinBox()
        self.spin_turn_interval.setValue(0.5)
        self.spin_turn_interval.setDecimals(3)
        label_turn_interval_unit = QLabel('秒')
        label_turn_interval_unit.setFixedWidth(12)
        btn_turn_test = QPushButton("单步调试")
        btn_turn_test.clicked.connect(self.one_step_test)
        # 1.2 控件导入布局
        self.grid.addWidget(label_page_turn, 4, 0, 1, 4)
        self.grid.addWidget(label_method, 5, 0)
        sublay_page_turn = QHBoxLayout()
        sublay_page_turn.addWidget(self.rad_method1)
        sublay_page_turn.addWidget(self.rad_method2)
        self.grid.addLayout(sublay_page_turn, 5, 1, 1, 3)
        self.grid.addWidget(label_mouse_pos0, 6, 0)
        self.grid.addWidget(self.spin_mouse_x0, 6, 1)
        self.grid.addWidget(self.spin_mouse_y0, 6, 2)
        self.grid.addWidget(btn_mouse_pos0_adjust, 6, 3)
        self.grid.addWidget(label_mouse_posp, 7, 0)
        self.grid.addWidget(self.spin_mouse_xp, 7, 1)
        self.grid.addWidget(self.spin_mouse_yp, 7, 2)
        self.grid.addWidget(btn_mouse_posp_adjust, 7, 3)
        self.grid.addWidget(label_page_ydisp, 8, 0)
        self.grid.addWidget(self.spin_page_ydisp, 8, 1, 1, 2)
        self.grid.addWidget(btn_page_ydisp_adjust, 8, 3)
        self.grid.addWidget(label_turn_interval, 9, 0)
        sublay_turn_interval = QHBoxLayout()
        sublay_turn_interval.addWidget(self.spin_turn_interval)
        sublay_turn_interval.addWidget(label_turn_interval_unit)
        self.grid.addLayout(sublay_turn_interval, 9, 1, 1, 2)
        self.grid.addWidget(btn_turn_test, 9, 3)

        # 2.输出控制
        # 2.1 创建控件
        label_output = QLabel('输出控制')
        label_output.setFont(self.font_title)
        label_page_start = QLabel('起始页码：')
        self.spin_page_start = self.create_spin(1, vmin = 1)
        label_page_count = QLabel('截取页数：')
        self.spin_page_count = self.create_spin(10, vmin = 1)
        label_pic = QLabel('保存图片：')
        self.dir_pic = QLineEdit()
        self.dir_pic.setText(self.cwd + '/pics')
        btn_sel_pic = QPushButton('…')
        btn_sel_pic.setMinimumWidth(35)
        btn_sel_pic.clicked.connect(self.sel_pic_path)
        btn_clear_pic = QPushButton('×')
        btn_clear_pic.setMinimumWidth(35)
        btn_clear_pic.clicked.connect(lambda: self.dir_pic.clear())
        label_pdf = QLabel('输出PDF：')
        self.dir_pdf = QLineEdit()
        btn_sel_pdf = QPushButton('…')
        btn_sel_pdf.setMinimumWidth(35)
        btn_sel_pdf.clicked.connect(self.sel_pdf_path)
        btn_clear_pdf = QPushButton('×')
        btn_clear_pdf.setMinimumWidth(35)
        btn_clear_pdf.clicked.connect(lambda: self.dir_pdf.clear())

        # 2.2 控件导入布局
        self.grid.addWidget(label_output, 10, 0, 1, 4)
        self.grid.addWidget(label_page_start, 11, 0)
        self.grid.addWidget(self.spin_page_start, 11, 1, 1, 2)
        self.grid.addWidget(label_page_count, 12, 0)
        self.grid.addWidget(self.spin_page_count, 12, 1, 1, 2)
        self.grid.addWidget(label_pic, 13, 0)
        self.grid.addWidget(self.dir_pic, 13, 1, 1, 2)
        sublay_pic = QHBoxLayout()
        sublay_pic.setSpacing(4)
        sublay_pic.addWidget(btn_sel_pic)
        sublay_pic.addWidget(btn_clear_pic)
        self.grid.addLayout(sublay_pic, 13, 3)
        self.grid.addWidget(label_pdf, 14, 0)
        self.grid.addWidget(self.dir_pdf, 14, 1, 1, 2)
        sublay_pdf = QHBoxLayout()
        sublay_pdf.setSpacing(4)
        sublay_pdf.addWidget(btn_sel_pdf)
        sublay_pdf.addWidget(btn_clear_pdf)
        self.grid.addLayout(sublay_pdf, 14, 3)
        
        # 4.截图按钮
        # 4.1 创建控件
        btn_test = QPushButton('测试翻页效果')
        btn_run = QPushButton('正式翻页截图')
        btn_test.clicked.connect(self.test)
        btn_run.clicked.connect(self.run)
        self.label_state = QLabel('【待运行】')
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        # 4.2 控件导入布局
        self.grid.addWidget(btn_test, 15, 0, 1, 2)
        self.grid.addWidget(btn_run, 15, 2, 1, 2)
        self.grid.addWidget(self.label_state, 16, 0)
        self.grid.addWidget(self.progress_bar, 16, 1, 1, 3)

    def update_method(self, method):
        if self.rad_method1.isChecked():
            self.method = 1
        else:
            self.method = 2

    def get_params(self):
        '''读取用户输入的参数值'''
        super(DocumentWindow, self).get_params()
        self.mouse_x0 = self.spin_mouse_x0.value()
        self.mouse_y0 = self.spin_mouse_y0.value()
        self.mouse_xp = self.spin_mouse_xp.value()
        self.mouse_yp = self.spin_mouse_yp.value()
        self.ydisp = self.spin_page_ydisp.value()
        self.turn_interval = self.spin_turn_interval.value()
        self.page_start = self.spin_page_start.value()
        self.page_count = self.spin_page_count.value()

    def get_ydisp(self):
        '''获得页向下位移选取对话框'''
        self.get_params()
        self.displacement_win = DisplacementDialog(self.sc_left, self.sc_right, self.mouse_x0, self.mouse_y0)
        self.displacement_win.setValueObj(self.spin_page_ydisp)
        self.displacement_win.show()

    def drag_page(self):
        '''垂直拖拽页面产生位移'''
        pyautogui.mouseDown(x = self.mouse_x0, y = self.mouse_y0, button = 'left')  #按下鼠标
        pyautogui.dragTo(x = self.mouse_x0, y = self.mouse_y0 - self.ydisp, duration = 0.2, button = 'left')  #鼠标拖拽
        pyautogui.mouseUp(x = self.mouse_x0, y = self.mouse_y0 - self.ydisp, button = 'left')  #松开鼠标

    def turn_page(self, page):
        '''进行翻页操作'''
        if self.method == 1:
            pyautogui.click(x = self.mouse_xp, y = self.mouse_yp, button = 'left', duration = 0.1)  #单击翻页键
        elif self.method == 2:
            pyautogui.doubleClick(x = self.mouse_xp, y = self.mouse_yp, duration = 0.1)  #鼠标挪到翻页框
            pyautogui.typewrite(message = str(page), interval = 0.1)  #写入新页码
            pyautogui.press('enter')  #按下回车翻页

    def one_step_test(self):
        '''自动翻页的单步测试'''
        self.get_params()
        x0, y0 = pyautogui.position()
        if self.ydisp != 0:
            self.drag_page()
        self.turn_page(self.page_start + 1)
        pyautogui.moveTo(x0, y0)

    def test(self):
        '''自动翻页测试'''
        self.get_params()
        self.label_state.setText('【测试中】')
        self.progress_bar.setValue(0)
        x0, y0 = pyautogui.position()
        for i in range(self.page_count):
            if self.ydisp != 0:
                self.drag_page()
                time.sleep(self.turn_interval)
            percent = int((i + 1) / self.page_count * 100)
            if percent > self.progress_bar.value():
                self.progress_bar.setValue(percent)
            self.turn_page(self.page_start + i + 1)
            time.sleep(self.turn_interval)
            QApplication.processEvents()
        pyautogui.moveTo(x0, y0)
        self.label_state.setText('【测试终】')

    def img_combine(self, f1, f2):
        '''将分别截取的两个图片拼接'''
        # 获得图像的像素点矩阵
        img1 = Image.open(f1)
        img2 = Image.open(f2)
        img1_info = img1.convert('L')
        img2_info = img2.convert('L')
        mat1 = np.array(img1_info)
        mat2 = np.array(img2_info)
        # 拿图片的最后四十行作比较
        mat1_tail = mat1[-40:,:]
        mat2_height = mat2.shape[0]
        # 获得相似度最高的行，以确定两张图片从何处拼接
        row_best = 0
        vbest = np.sum(np.abs(mat1_tail - mat2[:40,:]))
        for row in range(mat2_height - 41, -1, -1):
            v = np.sum(np.abs(mat1_tail - mat2[row:row+40,:]))
            if v < vbest:
                row_best = row
                vbest = v
        # 确定第2张图片截取范围
        img2_y0 = row_best + 40
        img2_height = mat2_height - img2_y0
        # 进行图像拼接
        comb_img = Image.new(img1.mode, (mat1.shape[1], mat1.shape[0] + img2_height))
        comb_img.paste(img1, box=(0,0))
        region_img2 = (0, img2_y0, mat2.shape[1], mat2.shape[0])
        comb_img.paste(img2.crop(region_img2), box=(0,mat1.shape[0]))
        # 保存图片
        comb_img.save(f1)

    def run(self):
        '''进行自动翻页及截图操作'''
        self.get_params()
        self.label_state.setText('【截图中】')
        self.progress_bar.setValue(0)
        x0, y0 = pyautogui.position()
        save_pic = self.dir_pic.text() != ''
        create_pdf = self.dir_pdf.text() != ''
        if not save_pic:
            fPath = self.cwd
        else:
            fPath = self.dir_pic.text()
        if create_pdf:
            self.init_pdf()
        for i in range(self.page_count):
            fName = '%s/p%03d.jpg' % (fPath, self.page_start + i)
            self.window_capture(fName)
            if self.ydisp != 0:
                self.drag_page()
                time.sleep(self.turn_interval)
                fName_part = '%s/p%03d_part.jpg' % (fPath, self.page_start + i)
                self.window_capture_part(fName_part, self.ydisp)
                self.img_combine(fName, fName_part)
                os.remove(fName_part)
            if create_pdf:
                self.append_pdf_page(fName)
            if not save_pic:
                os.remove(fName)
            percent = int((i + 1) / self.page_count * 100)
            if percent > self.progress_bar.value():
                self.progress_bar.setValue(percent)
            self.turn_page(self.page_start + i + 1)
            time.sleep(self.turn_interval)
            QApplication.processEvents()
        if create_pdf:
            self.end_pdf(self.dir_pdf.text())
        pyautogui.moveTo(x0, y0)
        self.label_state.setText('【截图终】')

if __name__ == '__main__':
    app = QApplication(sys.argv)
    doc_win = DocumentWindow()
    doc_win.show()
    sys.exit(app.exec_())